import glob
import multiprocessing
import os
import shutil
import sys
import tempfile
import time
import traceback

import histogram

from charmicat import utils



global _g_charmicat_mp_ncpu
_g_charmicat_mp_ncpu = 1

class MPMonitor(object):
    def __init__(self, nproc, logdir='./ccat_log', interval=5):

        if logdir is None:
            self.logdir = tempfile.mkdtemp(suffix='.ccat')
            self.my_cleanup = True
        else:
            self.logdir = utils.safe_mkdir(logdir, force_opt='--force-log')
            self.my_cleanup = False

        self.nproc = nproc
        self.interval = interval
        self.npolls = 0
        self.ndone = 0
        self.log_glob_pattern = '*.out'
        self.err_glob_pattern = '*.err'
        self.done_glob_pattern = '*.done'


    def __call__(self, *args):        
        self.npolls = 0
        time.sleep(self.interval)
        while True:
            errlogs = glob.glob(os.path.join(self.logdir, self.err_glob_pattern))
            self.npolls += 1

            if len(errlogs) > 0:
                for el in errlogs:
                    fp = open(el)
                    
                    msg = "Error encountered in process {0} after {1} polls ({2}/{3} jobs completed before error):\n"
                    sys.stderr.write(msg.format(os.path.basename(el).rsplit('.', 1)[0],
                                                self.npolls, self.ndone, self.nproc))
                    
                    sys.stderr.write("{0}\n".format(fp.read()))
                    sys.stderr.flush()
                    
                    fp.close()
                
                sys.stderr.write("Output left in `{0}'\n".format(self.logdir))
                sys.stderr.write("Aborting...\n")

                sys.exit()
            
            self.logfiles = glob.glob(os.path.join(self.logdir, self.log_glob_pattern))

            self.process_logs()
            
            self.ndone = len(glob.glob(os.path.join(self.logdir, self.done_glob_pattern)))

            if self.ndone >= self.nproc:
                break

            time.sleep(self.interval)


            
        self.finalize()            
        
        self.cleanup()

    def process_logs(self):
        pass

    def finalize(self):
        pass
    
    def cleanup(self):
        if self.my_cleanup:
            shutil.rmtree(self.logdir)


class CounterMonitor(MPMonitor):
    def process_logs(self):
        sys.stdout.write("{0}/{1} jobs complete (polls={2})\r".format(self.ndone,
                                                                      self.nproc,
                                                                      self.npolls))

        sys.stdout.flush()

    
    def finalize(self):
        print "{0}/{1} jobs complete (polls={2})".format(self.ndone,
                                                         self.nproc,
                                                         self.npolls)

class FillMonitor(MPMonitor):
    def __init__(self, nproc, nfiles, logdir, opts):
        super(FillMonitor, self).__init__(nproc, logdir=logdir, interval=2)
        
        self.cmdline_opts = opts
        self.nfiles_tot   = nfiles
        
        self.hcoll = histogram.HistogramCollection()

    def process_logs(self):
        nfiles_all_procs = 0
        for path in self.logfiles:
            fp = open(path)
            
            lines = [line for line in fp if line[0] != '#']
            try:
                nfile = lines[-1].split()[0]
            except IndexError:
                nfile = 0
            except ValueError:
                nfile = 0
            
            nfiles_all_procs += int(nfile)
            
            fp.close()

        
        self.notify(nfiles_all_procs)

    def notify(self, nfiles_all_procs):
        sys.stdout.write("{0}/{1} files complete (polls={2})\r".format(nfiles_all_procs, self.nfiles_tot, self.npolls))
        sys.stdout.flush()        
        
    def finalize(self):
        print
        hpt_paths = glob.glob(os.path.join(self.logdir, '*.hpt'))

        for path in hpt_paths:
            tmp_hcoll = histogram.load(path)

            try:
                iter(tmp_hcoll)
            except TypeError:
                tmp_hcoll = [tmp_hcoll]
            
            for h in tmp_hcoll:
                if h.name in self.hcoll:
                    self.hcoll[h.name] += h
                else:
                    self.hcoll.add(h)
                



class MPWorker(object):
    def __init__(self, logdir, args=[], logfile_tag=None, errfile_tag=None, donefile_tag=None):
        self.logdir = logdir
        self.args = args
                
        self.logfile = None
        self.errlog  = None

        self.logfile_tag  = logfile_tag
        self.errfile_tag  = errfile_tag
        self.donefile_tag = donefile_tag


    def __call__(self):
        if self.logfile_tag:
            self.logpath = os.path.join(self.logdir, '{0}.{1}.out'.format(os.getpid(), self.logfile_tag))
        else:
            self.logpath = os.path.join(self.logdir, '{0}.out'.format(os.getpid()))

        if self.errfile_tag:
            self.errpath = os.path.join(self.logdir, '{0}.{1}.err'.format(os.getpid(), self.errfile_tag))
        else:
            self.errpath = os.path.join(self.logdir, '{0}.err'.format(os.getpid()))
        
            
        try:
            self.work(*self.args)
        except Exception, e:
            self.error(traceback.format_exc())
            self.cleanup()
            return

        self.completed()

    def cleanup(self):
        if self.logfile:
            self.logfile.close()
        if self.errlog:
            self.errlog.close()            

    def work(self, *args):
        pass
    
    def completed(self):
        if self.donefile_tag is not None:            
            open(os.path.join(self.logdir, '{0}.{1}.done'.format(os.getpid(), self.donefile_tag)), 'w').close()
        else:
            open(os.path.join(self.logdir, '{0}.done'.format(os.getpid())), 'w').close()

    def log(self, *args):
        self.logfile = open(self.logpath, 'a+')        
        self.logfile.write(" ".join(map(str, args)) + '\n')
        self.logfile.flush()

    def error(self, *args):
        self.errlog = open(self.errpath, 'a+')
        self.errlog.write(" ".join(map(str, args)) + '\n')
        self.errlog.close()

def mp_dummy():
    pass

def mp_default_callback(arg):
    return arg

def mp_deploy_and_monitor(monitor, workers, ncpu=None, cbs=None):
    global _g_charmicat_mp_ncpu
    if ncpu is not None:
        if ncpu <= multiprocessing.cpu_count():
            _g_charmicat_mp_ncpu = ncpu
        else:
            _g_charmicat_mp_ncpu = multiprocessing.cpu_count()
    else:
        _g_charmicat_mp_ncpu = multiprocessing.cpu_count()
        
    if _g_charmicat_mp_ncpu == 1:
        for worker in workers:
            worker()

        monitor()
    else:
        thread_pool = multiprocessing.Pool(_g_charmicat_mp_ncpu)
        procs = [mp_dummy] + workers
        if cbs is None:
            cbs   = [monitor] + [mp_default_callback] * len(workers)
            
        for ip, proc in enumerate(procs):
            thread_pool.apply_async(proc, callback=cbs[ip])

        thread_pool.close()

        thread_pool.join()
